"""
RAG Embeddings - 文档处理模块

该模块负责PDF文档的文本提取和文本分块功能：
- 从PDF文件中提取文本内容
- 将长文本按指定大小进行分块处理
- 支持分块重叠以保持上下文连贯性

Author: xiaodeme@163.com
Date: 2025-07-11 10:24:41
"""
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer

from config import Config


class Document:
    """文档处理类，用于PDF文本提取和分块"""
    
    def __init__(self, pdf_path=None, page_numbers=None):
        """
        初始化文档处理器
        
        Args:
            pdf_path (str, optional): PDF文件路径
            page_numbers (list, optional): 指定要处理的页码列表，None表示处理所有页面
        """
        self.pdf_path = pdf_path
        self.page_numbers = page_numbers

    def extract_text_from_pdf(self):
        """
        从PDF文件中提取文本内容
        
        该方法会逐页扫描PDF文件，提取所有文本内容并合并为一个完整的字符串，
        自动去除所有空白字符（包括空格、换行、制表符等）
        
        Returns:
            str: 提取的完整文本内容（无空白字符）
        """
        if not self.pdf_path:
            print("[ERROR] PDF文件路径未设置")
            return ""
            
        full_text = ""
        
        try:
            # 逐页处理PDF文件
            for page_index, page_layout in enumerate(extract_pages(self.pdf_path)):
                # 如果指定了页码范围，跳过不在范围内的页面
                if self.page_numbers is not None and page_index not in self.page_numbers:
                    continue

                # 提取当前页面中的所有文本元素
                for element in page_layout:
                    if isinstance(element, LTTextContainer):
                        # 获取文本块内容并去除首尾空白
                        text_block = element.get_text().strip()
                        # 添加到完整文本中
                        full_text += text_block

            # 删除字符串中所有的空白字符（包括空格、换行、制表符等）
            full_text = ''.join(full_text.split())
            
            print(f"[INFO] 成功从PDF提取文本，总字符数: {len(full_text)}")
            return full_text

        except Exception as e:
            print(f"[ERROR] PDF文本提取失败: {e}")
            return ""

    def split_text(self, contents=None, chunk_size=Config.CHUNK_SIZE, chunk_overlap=Config.CHUNK_OVERLAP):
        """
        将完整文本按指定大小进行分块处理
        
        该方法支持设置块大小和重叠字符数，重叠部分有助于保持上下文的连贯性
        
        Args:
            contents (str): 待分块的完整文本字符串
            chunk_size (int): 每个文本块的最大字符数，默认512
            chunk_overlap (int): 相邻文本块之间的重叠字符数，默认20

        Returns:
            list: 切分后的文本块列表（每个元素为字符串）
        """
        if not contents or not isinstance(contents, str):
            print("[WARNING] 输入文本为空或格式错误")
            return []

        chunks = []
        start = 0

        while start < len(contents):
            # 计算当前块的结束位置
            end = start + chunk_size
            chunk = contents[start:end]
            chunks.append(chunk)
            
            # 计算下一个块的起始位置（考虑重叠）
            start += chunk_size - chunk_overlap

        print(f"[INFO] 文本分块完成，共生成 {len(chunks)} 个文本块")
        return chunks


if __name__ == '__main__':
    """测试代码：演示PDF文本提取和分块功能"""
    
    # 使用配置文件中的PDF路径
    pdf_path = Config.PDF_PATH 
    document = Document(pdf_path)
    
    # 提取PDF文本内容
    contents = document.extract_text_from_pdf()
    print(f"提取的PDF完整内容长度: {len(contents)} 字符")

    # 将文本进行分块处理
    chunks = document.split_text(contents)
    for i, chunk in enumerate(chunks):
        print(f">>> 分块 {i+1} >>> (长度: {len(chunk)}): {chunk[:50]}...")  # 只显示前50个字符作为预览

    